from numpy import array
from sklearn.externals import joblib
from sklearn.naive_bayes import MultinomialNB
from get_words_from_file import getWordsFromFile

model = joblib.load("垃圾邮件分类器.pkl")
print('加载模型和训练结果成功。')
with open('topWords.txt', encoding='utf8') as fp:
    topWords = fp.read().split(',')

def predict(txtFile):
    # 获取指定邮件文件内容，返回分类结果
    words = getWordsFromFile(txtFile)
    currentVector = array(tuple(map(lambda x: words.count(x),
                                    topWords)))
    result = model.predict(currentVector.reshape(1, -1))[0]
    return '垃圾邮件' if result==1 else '正常邮件'

# 151.txt至155.txt为测试邮件内容
for mail in ('%d.txt'%i for i in range(151, 156)):
    print(mail, predict(mail), sep=':')
